package com.lmabbe.oauth.config;

import com.lmabbe.common.global.context.session.JwtUser;
import com.lmabbe.common.global.result.R;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.MethodParameter;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.security.authentication.*;
import org.springframework.security.core.userdetails.UsernameNotFoundException;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.exceptions.OAuth2Exception;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.error.WebResponseExceptionTranslator;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.util.HashMap;
import java.util.Map;

/**
 * @author lmabbe
 * @data 2021/5/16 20:53
 */
@Slf4j
@ControllerAdvice
public class MyResponseBodyAdvice implements ResponseBodyAdvice, WebResponseExceptionTranslator {

    @Autowired
    private TokenStore tokenStore;

    @Override
    public boolean supports(MethodParameter methodParameter, Class aClass) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter methodParameter, MediaType mediaType, Class aClass, ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse) {
        log.info("请求返回数据类型class=" + body.getClass().getName());
        if (body instanceof DefaultOAuth2AccessToken) {
            DefaultOAuth2AccessToken defaultOAuth2AccessToken = (DefaultOAuth2AccessToken) body;
            OAuth2Authentication oAuth2Authentication = tokenStore.readAuthentication(defaultOAuth2AccessToken.getValue());
            JwtUser jwtUser = (JwtUser) oAuth2Authentication.getPrincipal();
            Map<String, Object> extra = new HashMap<>();
            extra.put("token", defaultOAuth2AccessToken.getValue());
            extra.put("refreshToken", defaultOAuth2AccessToken.getRefreshToken().getValue());
            extra.put("expiration", defaultOAuth2AccessToken.getExpiration());
            return R.success(jwtUser.getLoginUser()).setExtra(extra);
        }
        return body;
    }

    @Override
    public ResponseEntity translate(Exception exception) throws Exception {
        exception.printStackTrace();
        log.info("异常类型：" + exception.getClass().getName());
        if (exception instanceof OAuth2Exception) {
            return new ResponseEntity(R.fail("认证异常：" + exception.getMessage()), HttpStatus.OK);
        } else if (exception instanceof LockedException) {
            return new ResponseEntity(R.fail("账户被锁定，登陆失败！"), HttpStatus.OK);
        } else if (exception instanceof BadCredentialsException ||
                exception instanceof UsernameNotFoundException) {
            return new ResponseEntity(R.fail("账户或者密码错误，登陆失败！"), HttpStatus.OK);
        } else if (exception instanceof DisabledException) {
            return new ResponseEntity(R.fail("账户被禁用，登陆失败！"), HttpStatus.OK);
        } else if (exception instanceof AccountExpiredException) {
            return new ResponseEntity(R.fail("账户已过期，登陆失败！"), HttpStatus.OK);
        } else if (exception instanceof CredentialsExpiredException) {
            return new ResponseEntity(R.fail("密码已过期，登陆失败！"), HttpStatus.OK);
        } else {
            return new ResponseEntity(R.fail("未知错误:" + exception.getMessage()), HttpStatus.OK);
        }
    }
}
